import numpy as np
import pandas as pd
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import accuracy_score
import matplotlib.pyplot as plt
import seaborn as sns
from collections import Counter

# 加载数据集
iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)


# 自定义KNN类
class KNN:
    def __init__(self, k=3, distance_metric='euclidean'):
        self.k = k
        self.distance_metric = distance_metric

    def fit(self, X, y):
        self.X_train = X
        self.y_train = y

    def predict(self, X):
        y_pred = [self._predict(x) for x in X]
        return np.array(y_pred)

    def _predict(self, x):
        # 计算距离
        if self.distance_metric == 'euclidean':
            distances = [self._euclidean_distance(x, x_train) for x_train in self.X_train]
        elif self.distance_metric == 'manhattan':
            distances = [self._manhattan_distance(x, x_train) for x_train in self.X_train]
        elif self.distance_metric == 'minkowski':
            distances = [self._minkowski_distance(x, x_train) for x_train in self.X_train]
        else:
            raise ValueError("不支持的距离度量方法")

        # 获取最近的k个样本的索引
        k_indices = np.argsort(distances)[:self.k]
        # 获取这些样本的标签
        k_nearest_labels = [self.y_train[i] for i in k_indices]
        # 投票决定预测标签
        most_common = Counter(k_nearest_labels).most_common(1)
        return most_common[0][0]

    def _euclidean_distance(self, x1, x2):
        return np.sqrt(np.sum((x1 - x2) ** 2))

    def _manhattan_distance(self, x1, x2):
        return np.sum(np.abs(x1 - x2))

    def _minkowski_distance(self, x1, x2, p=3):
        return np.power(np.sum(np.abs(x1 - x2) ** p), 1 / p)


# 寻找最优K值
def find_best_k(X_train, y_train, X_test, y_test, distance_metric='euclidean'):
    accuracies = []
    k_range = range(1, 31)

    for k in k_range:
        knn = KNN(k=k, distance_metric=distance_metric)
        knn.fit(X_train, y_train)
        y_pred = knn.predict(X_test)
        accuracy = accuracy_score(y_test, y_pred)
        accuracies.append(accuracy)

    best_k = k_range[np.argmax(accuracies)]
    best_accuracy = max(accuracies)

    plt.figure(figsize=(10, 6))
    plt.plot(k_range, accuracies, marker='o')
    plt.title(f'不同K值的准确率 ({distance_metric}距离)')
    plt.xlabel('K值')
    plt.ylabel('准确率')
    plt.xticks(k_range)
    plt.grid(True)
    plt.savefig(f'k_accuracy_{distance_metric}.png')
    plt.show()

    return best_k, best_accuracy


# 使用不同距离度量方法寻找最优K值
distance_metrics = ['euclidean', 'manhattan', 'minkowski']
results = {}

for metric in distance_metrics:
    print(f"\n使用{metric}距离度量方法:")
    best_k, best_accuracy = find_best_k(X_train, y_train, X_test, y_test, metric)
    results[metric] = {'best_k': best_k, 'accuracy': best_accuracy}
    print(f"最优K值: {best_k}")
    print(f"最高准确率: {best_accuracy:.4f}")

# 找出所有距离度量方法中最好的结果
best_metric = max(results, key=lambda x: results[x]['accuracy'])
print(f"\n所有距离度量方法中最优的结果:")
print(f"距离度量: {best_metric}")
print(f"最优K值: {results[best_metric]['best_k']}")
print(f"最高准确率: {results[best_metric]['accuracy']:.4f}")

# 使用sklearn的KNN进行对比
best_k_sklearn = results[best_metric]['best_k']
knn_sklearn = KNeighborsClassifier(n_neighbors=best_k_sklearn, metric=best_metric)
knn_sklearn.fit(X_train, y_train)
y_pred_sklearn = knn_sklearn.predict(X_test)
accuracy_sklearn = accuracy_score(y_test, y_pred_sklearn)

print(f"\nsklearn的KNN结果 (K={best_k_sklearn}, 距离={best_metric}):")
print(f"准确率: {accuracy_sklearn:.4f}")

# 比较自己实现的KNN和sklearn的KNN结果
if 'best_k' in results[best_metric]:
    best_k = results[best_metric]['best_k']
    knn_custom = KNN(k=best_k, distance_metric=best_metric)
    knn_custom.fit(X_train, y_train)
    y_pred_custom = knn_custom.predict(X_test)

    # 创建混淆矩阵
    from sklearn.metrics import confusion_matrix

    cm_custom = confusion_matrix(y_test, y_pred_custom)
    cm_sklearn = confusion_matrix(y_test, y_pred_sklearn)

    # 可视化混淆矩阵
    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    sns.heatmap(cm_custom, annot=True, fmt='d', cmap='Blues',
                xticklabels=iris.target_names, yticklabels=iris.target_names)
    plt.title(f'自定义KNN混淆矩阵 (K={best_k}, {best_metric})')
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')

    plt.subplot(1, 2, 2)
    sns.heatmap(cm_sklearn, annot=True, fmt='d', cmap='Greens',
                xticklabels=iris.target_names, yticklabels=iris.target_names)
    plt.title(f'sklearn KNN混淆矩阵 (K={best_k}, {best_metric})')
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')

    plt.tight_layout()
    plt.savefig('confusion_matrices.png')
    plt.show()
